package at.ltasat.classifier.dt;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import at.ltasat.classifier.Classifier;
import at.ltasat.classifier.ClassifierFactory;
import at.ltasat.classifier.ClazzManager;
import at.ltasat.classifier.Membership;
import at.ltasat.classifier.ObjectFeaturesFactory;
import at.ltasat.classifier.Trainer;
import at.ltasat.util.IO;

public final class DecisionTreeClassifierFactory<O, D> implements ClassifierFactory<O, D> {

	public static final String NAME = "dt";
	private static final byte DECISION_NODE_HEADER = 0x1;
	private static final byte CLASSIFICATION_NODE_HEADER = 0x2;

	@Override
	public String getName() {
		return NAME;
	}

	@Override
	public Trainer<O, D> createTrainer(ClazzManager<D> classes, ObjectFeaturesFactory<O> factory) {
		return new DecisionTreeTrainer<O, D>(NAME, classes, factory);
	}

	@Override
	public Classifier<O, D> readClassifier(ClazzManager<D> classes, ObjectFeaturesFactory<O> factory, InputStream in) throws IOException {
		return new DecisionTreeClassifier<O, D>(NAME, classes, factory, readNode(in, classes));
	}

	public static <O, D> void writeClassifier(final OutputStream out, final DecisionTreeClassifier<O, D> classifier) throws IOException {
		writeNode(out, classifier.getRoot());
	}

	private static <D> void writeNode(final OutputStream out, final Node<D> n) throws IOException {
		if (n instanceof DecisionNode<?>) {
			DecisionNode<D> node = (DecisionNode<D>) n;
			IO.writeByte(out, DECISION_NODE_HEADER);
			IO.writeInteger(out, node.getFeature());
			IO.writeDouble(out, node.getSplit());
			writeNode(out, node.getLower());
			writeNode(out, node.getUpper());
		} else if (n instanceof ClassificationNode<?>) {
			ClassificationNode<D> node = (ClassificationNode<D>) n;
			IO.writeByte(out, CLASSIFICATION_NODE_HEADER);
			IO.writeInteger(out, node.getNumberOfMemberships());
			for (int i = 0; i < node.getNumberOfMemberships(); i++) {
				final Membership<D> membership = node.getMembership(i);
				IO.writeInteger(out, membership.getClazz().getId());
				IO.writeDouble(out, membership.getDegree());
			}
		} else {
			assert false;
		}
	}

	@SuppressWarnings("unchecked")
	private static <D> Node<D> readNode(final InputStream in, final ClazzManager<D> clazzes) throws IOException {
		Node<D> node;
		final byte header = IO.readByte(in);
		switch (header) {
		case DECISION_NODE_HEADER:
			final int feature = IO.readInteger(in);
			final double split = IO.readDouble(in);
			final Node<D> lower = readNode(in, clazzes);
			final Node<D> upper = readNode(in, clazzes);
			node = new DecisionNode<D>(feature, split, lower, upper);
			break;
		case CLASSIFICATION_NODE_HEADER:
			final Membership<D>[] memberships = (Membership<D>[]) new Membership<?>[IO.readInteger(in)];
			for (int count = 0; count < memberships.length; count++) {
				final int clazzId = IO.readInteger(in);
				final double degree = IO.readDouble(in);
				memberships[count] = new Membership<D>(clazzes.get(clazzId), degree);
			}
			node = new ClassificationNode<D>(memberships);
			break;
		default:
			assert false;
			node = null;
			break;
		}
		return node;
	}

}
